const std = @import("std");

const VLEN = 32;
const VECF = @Vector(VLEN, f64);
const VECI = @Vector(VLEN, u32);

fn eval_A(i_param: VECI, j_param: VECI) VECF {
    const fone: VECF = @splat(1.0);
    const ione: VECI = @splat(1);
    const itwo: VECI = @splat(2);
    //     return 1.0 / @as(f64, @floatFromInt((i_param + j_param) * (i_param + j_param + 1) / 2 + i_param + 1));
    const iv: VECI = (i_param + j_param) * (i_param + j_param + ione) / itwo + i_param + ione;
    const fv: VECF = @floatFromInt(iv);
    return fone / fv;
}

fn eval_A_times_u(us: []const VECF, Aus: []VECF) void {
    for (Aus, 0..) |*Au, ii| {
        const ni = ii * VLEN;
        const i: VECI = blk: {
            var vs: [VLEN]u32 = undefined;
            for (&vs, 0..) |*v, k| {
                v.* = @intCast(ni + k);
            } else {
                break :blk vs[0..].*;
            }
        };

        Au.* = @splat(0);
        for (us, 0..) |uu, jj| {
            const nj = jj * VLEN;
            for (0..VLEN) |h| {
                const j: VECI = @splat(@intCast(nj + h));
                const u: VECF = @splat(uu[h]);

                Au.* += eval_A(i, j) * u;
            }
        }
    }

    //     for (0..N) |i| {
    //         Au[i] = 0.0;
    //         for (0..N) |j| {
    //             Au[i] += eval_A(@intCast(i), @intCast(j)) * u[j];
    //         }
    //     }
}

fn eval_At_times_u(us: []const VECF, Aus: []VECF) void {
    for (Aus, 0..) |*Au, ii| {
        const ni = ii * VLEN;
        const i: VECI = blk: {
            var vs: [VLEN]u32 = undefined;
            for (&vs, 0..) |*v, k| {
                v.* = @intCast(ni + k);
            } else {
                break :blk vs[0..].*;
            }
        };

        Au.* = @splat(0);
        for (us, 0..) |uu, jj| {
            const nj = jj * VLEN;
            for (0..VLEN) |h| {
                const j: VECI = @splat(@intCast(nj + h));
                const u: VECF = @splat(uu[h]);

                Au.* += eval_A(j, i) * u;
            }
        }
    }

    //     for (0..N) |i| {
    //         Au[i] = 0.0;
    //         for (0..N) |j| {
    //             // Note the swapped arguments to eval_A: (j, i) instead of (i, j)
    //             Au[i] += eval_A(@intCast(j), @intCast(i)) * u[j];
    //         }
    //     }
}

fn eval_AtA_times_u(u: []const VECF, AtAu: []VECF, temp_v_buf: []VECF) void {
    eval_A_times_u(u, temp_v_buf);
    eval_At_times_u(temp_v_buf, AtAu);
}

pub fn main() !void {
    @setFloatMode(.optimized);
    const stdout = std.io.getStdOut().writer();
    const allocator = std.heap.smp_allocator;

    const args_list = try std.process.argsAlloc(allocator);
    defer std.process.argsFree(allocator, args_list);

    const N: usize = if (args_list.len == 2)
        try std.fmt.parseInt(usize, args_list[1], 10)
    else
        5472;

    const VN = N / VLEN;
    if (N % VLEN != 0) {
        try stdout.print("Input number must be a multiple of {d}\n", .{VLEN});
        return error.InputError;
    }

    const u_arr = try allocator.alloc(VECF, VN);
    defer allocator.free(u_arr);

    const v_arr = try allocator.alloc(VECF, VN);
    defer allocator.free(v_arr);

    const temp_v_for_AtA = try allocator.alloc(VECF, VN);
    defer allocator.free(temp_v_for_AtA);

    for (u_arr) |*u| {
        u.* = @splat(1);
    }

    var i_loop: usize = 0;
    while (i_loop < 10) : (i_loop += 1) {
        eval_AtA_times_u(u_arr, v_arr, temp_v_for_AtA);
        eval_AtA_times_u(v_arr, u_arr, temp_v_for_AtA);
    }

    //     var vBv: f64 = 0.0;
    //     var vv: f64 = 0.0;
    //
    //     for (0..N) |idx| {
    //         vBv += u_arr[idx] * v_arr[idx];
    //         vv += v_arr[idx] * v_arr[idx];
    //     }
    //     const result = std.math.sqrt(vBv / vv);

    var vBv: VECF = @splat(0);
    var vv: VECF = @splat(0);
    for (u_arr, v_arr) |u, v| {
        vBv += u * v;
        vv += v * v;
    }
    const result = std.math.sqrt(@reduce(.Add, vBv) / @reduce(.Add, vv));

    try stdout.print("{d:.9}\n", .{result});
}
